import torch
import pytorch_lightning as pl
import torch.nn.functional as F
import importlib
import torch.nn as nn

from .Encoder import Encoder
from .decoder import SwinIR
from .DiagonalGaussianDistribution import DiagonalGaussianDistribution

# from Encoder import Encoder
# from decoder import SwinIR
# from DiagonalGaussianDistribution import DiagonalGaussianDistribution
# import losses

class VAE(nn.Module):
    def __init__(self, embed_dim=3):
        super().__init__()
        self.encoder = Encoder(double_z=False, z_channels=3, resolution=160, in_channels=6, out_ch=3, ch=128,
                               ch_mult=[1, 2, 4], num_res_blocks=2, attn_resolutions=[])
        self.decoder = SwinIR(upscale=4, img_size=(40, 40),
                              window_size=8, img_range=1., depths=[6, 6, 6, 6],
                              embed_dim=96, num_heads=[6, 6, 6, 6], mlp_ratio=2, upsampler='pixelshuffle')
        self.quant_conv = torch.nn.Conv2d(3, 2 * embed_dim, 1)
        self.post_quant_conv = torch.nn.Conv2d(embed_dim, 3, 1)

    def encode(self, x):
        h = self.encoder(x)
        moments = self.quant_conv(h)
        posterior = DiagonalGaussianDistribution(moments)
        return posterior

    def decode(self, y, z):
        z = self.post_quant_conv(z)
        dec = self.decoder(y, z)
        return dec

    def get_last_layer(self):
        return self.decoder.conv_last.weight

    def forward(self, input, lr, sample_posterior=True):
        posterior = self.encode(input)
        if sample_posterior:
            z = posterior.sample()
        else:
            z = posterior.mode()
        dec = self.decode(lr, z)
        return dec, posterior

# vae = VAE(3)
# # print(vae.decoder)
# x = torch.randn((1, 6, 160, 160))
# lr = torch.randn((1, 3, 40, 40))
# # y = vae.encode(x)
# # print(y.shape)
# y, z = vae(x, lr)
# print(y.shape, z.shape)
